import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, input_sizes):
        super().__init__()
        self.fcs = nn.ModuleList()
        for i in range(len(input_sizes) - 1):
            if i != len(input_sizes) - 2:
                self.fcs.append(
                    nn.Linear(input_sizes[i], input_sizes[i + 1])
                )
                self.fcs.append(
                    nn.ReLU()
                )
            else:
                self.fcs.append(
                    nn.Linear(input_sizes[i], input_sizes[i + 1]),
                )

    def forward(self, x):
        for fc in self.fcs:
            x = fc(x)
        return x


if __name__ == '__main__':
    input_sizes = [1, 128, 256, 128, 64, 1]
    model = Model(input_sizes)
    # 1. 所有参数的生成器
    params = model.parameters()
    for param in params:
        print(param.shape)
    # 2. 获取所有参数和名称
    for name, param in model.named_parameters():
        print(f"name: {name}, param.shape:{param.shape}")
    # 3. 获取所有子模块
    for m in model.modules():
        print(m)
    # 4. 获取模块名称和模块本身
    for name, m in model.named_modules():
        print(f"name: {name}, module: {m}")